import gradio as gr
from Core.comfy_client import ComfyUIClient
import json
import os
import glob
import tempfile
import numpy as np
from PIL import Image
import cv2
import time
import uuid
import requests
import websocket

class BatchUpscaleInterface:
    def __init__(self):
        self.client = ComfyUIClient()
        self.workflow_path = os.path.join(os.path.dirname(__file__), "..", "Comfyui", "work_flow", "（批量）图片清晰放大1.0.json")
        
        # 获取模型路径
        from config.config import get_checkpoint_dir_path
        self.checkpoint_dir = get_checkpoint_dir_path()
        self.temp_dir = os.path.join(os.path.dirname(__file__), "..", "temp")
        os.makedirs(self.temp_dir, exist_ok=True)
        
        # 获取放大模型列表
        self.upscale_models = []
        upscale_model_dir = os.path.join(os.path.dirname(__file__), "..", "Comfyui", "models", "upscale_models")
        if os.path.exists(upscale_model_dir):
            for file in os.listdir(upscale_model_dir):
                if file.lower().endswith(('.pth', '.pt')):
                    self.upscale_models.append(file)
        
        # 如果没有找到模型，使用默认模型
        if not self.upscale_models:
            self.upscale_models = ["RealESRGAN_x4plus.pth"]
        
        # 获取工作流中的默认模型
        self.default_upscale_model = None
        try:
            workflow = self.client.load_workflow(self.workflow_path)
            if "53" in workflow and "inputs" in workflow["53"] and "model_name" in workflow["53"]["inputs"]:
                self.default_upscale_model = workflow["53"]["inputs"]["model_name"]
        except Exception as e:
            print(f"Error loading workflow for default upscale model: {str(e)}")
    
    def count_images_in_folder(self, folder_path):
        """统计文件夹中图片文件的数量"""
        if not folder_path or not os.path.exists(folder_path):
            return 0
        
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
        count = 0
        
        for file in os.listdir(folder_path):
            if any(file.lower().endswith(ext) for ext in image_extensions):
                count += 1
        
        return count
    
    def load_workflow(self):
        """加载工作流文件"""
        try:
            with open(self.workflow_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"Error loading workflow: {str(e)}")
            return None
    
    def get_image_files(self, folder_path):
        """获取文件夹中所有图片文件的路径"""
        if not folder_path or not os.path.exists(folder_path):
            return []
        
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
        image_files = []
        
        for file in os.listdir(folder_path):
            if any(file.lower().endswith(ext) for ext in image_extensions):
                image_files.append(os.path.join(folder_path, file))
        
        return sorted(image_files)
    
    def convert_ui_workflow_to_api_format(self, ui_workflow):
        """将UI导出的工作流转换为ComfyUI API格式"""
        api_workflow = {}
        
        for node in ui_workflow.get('nodes', []):
            node_id = str(node['id'])
            api_workflow[node_id] = {
                "class_type": node['type'],
                "inputs": {}
            }
            
            # 处理输入连接
            for input_info in node.get('inputs', []):
                if input_info.get('link') is not None:
                    # 找到连接的源节点
                    link_id = input_info['link']
                    for link in ui_workflow.get('links', []):
                        if link[0] == link_id:
                            source_node_id = str(link[1])
                            source_output_index = link[2]
                            api_workflow[node_id]["inputs"][input_info['name']] = [source_node_id, source_output_index]
                            break
            
            # 处理widget值 - 根据节点类型和实际的widgets_values设置
            if 'widgets_values' in node and node['widgets_values']:
                widget_values = node['widgets_values']
                
                if node['type'] == 'Load Image Batch':
                    # Load Image Batch的widgets_values顺序:
                    # [0]: mode, [1]: seed, [2]: action, [3]: batch_index, [4]: label, [5]: path, [6]: pattern, [7]: allow_RGBA_output, [8]: filename_text_extension
                    if len(widget_values) > 8:
                        api_workflow[node_id]["inputs"]["mode"] = widget_values[0]
                        api_workflow[node_id]["inputs"]["seed"] = widget_values[1] 
                        api_workflow[node_id]["inputs"]["index"] = widget_values[3]  # 这是batch_index，对应index输入
                        api_workflow[node_id]["inputs"]["label"] = widget_values[4]
                        api_workflow[node_id]["inputs"]["path"] = widget_values[5]  # 这个会被动态替换
                        api_workflow[node_id]["inputs"]["pattern"] = widget_values[6]
                        api_workflow[node_id]["inputs"]["allow_RGBA_output"] = widget_values[7]
                        api_workflow[node_id]["inputs"]["filename_text_extension"] = widget_values[8]
                        
                elif node['type'] == 'UpscaleModelLoader':
                    # UpscaleModelLoader只有一个model_name参数
                    api_workflow[node_id]["inputs"]["model_name"] = widget_values[0]
                    
                elif node['type'] == 'SaveImage':
                    # SaveImage只有一个filename_prefix参数
                    api_workflow[node_id]["inputs"]["filename_prefix"] = widget_values[0]
                    
                elif node['type'] == 'ImageScale':
                    # ImageScale的参数
                    if len(widget_values) >= 4:
                        api_workflow[node_id]["inputs"]["upscale_method"] = widget_values[0]
                        api_workflow[node_id]["inputs"]["width"] = widget_values[1]
                        api_workflow[node_id]["inputs"]["height"] = widget_values[2]
                        api_workflow[node_id]["inputs"]["crop"] = widget_values[3]
        
        return api_workflow

    def batch_upscale_images(self, folder_path, upscale_model, output_prefix, progress=gr.Progress()):
        """按文件数量循环执行工作流进行图片放大"""
        try:
            if not folder_path or not os.path.exists(folder_path):
                return "错误：请提供有效的文件夹路径", []

            # 获取图片文件列表并统计数量
            image_files = self.get_image_files(folder_path)
            total = len(image_files)
            if total == 0:
                return "错误：文件夹中没有找到图片文件", []

            # 加载UI工作流
            ui_workflow = self.load_workflow()
            if not ui_workflow:
                return "错误：无法加载工作流文件", []

            # 转换为API格式
            workflow = self.convert_ui_workflow_to_api_format(ui_workflow)

            # 设置批量加载路径 (Node 51)
            if "51" in workflow:
                workflow["51"]["inputs"]["path"] = folder_path.replace("\\", "/")
                # 确保模式和pattern等必要入参存在
                workflow["51"]["inputs"].setdefault("mode", "incremental_image")
                workflow["51"]["inputs"].setdefault("pattern", "*")
                print(f"设置批量加载路径: {folder_path}")

            # 设置放大模型 (Node 53)
            if "53" in workflow and upscale_model:
                workflow["53"]["inputs"]["model_name"] = upscale_model
                print(f"设置放大模型: {upscale_model}")

            # 设置输出前缀 (Node 56)
            if "56" in workflow and output_prefix:
                workflow["56"]["inputs"]["filename_prefix"] = output_prefix
                print(f"设置输出前缀: {output_prefix}")

            print("开始批量处理...")
            print(f"API工作流: {json.dumps(workflow, indent=2, ensure_ascii=False)}")

            results = []
            progress(0, desc="开始处理...")

            # 将轮询超时时间设置为5分钟
            retry_interval = 2  # 秒
            max_retries = 150   # 150 * 2s = 300s

            # 循环执行：index = 0..total-1
            for i in range(total):
                # 设置当前批量索引 (Node 51)
                if "51" in workflow:
                    workflow["51"]["inputs"]["index"] = i

                try:
                    resp = self.client.post_prompt(workflow)
                    print(f"API response: {resp}")
                    if "error" in resp:
                        return f"ComfyUI错误: {resp['error']}", []
                    prompt_id = resp["prompt_id"]

                    # 轮询等待输出，最长5分钟
                    for retry in range(max_retries):
                        hist = self.client.get_history(prompt_id)
                        if prompt_id in hist and "outputs" in hist[prompt_id]:
                            out = hist[prompt_id]["outputs"]
                            for node_id in out:
                                if "images" in out[node_id]:
                                    for img in out[node_id]["images"]:
                                        img_bytes = self.client.get_image(img["filename"], img["subfolder"], img["type"])
                                        file_path = os.path.join(self.temp_dir, f"upscaled_{prompt_id}_{i}_{img['filename']}")
                                        with open(file_path, "wb") as f:
                                            f.write(img_bytes)
                                        results.append(os.path.abspath(file_path))
                            break

                        time.sleep(retry_interval)
                        progress(((i + retry / max_retries) / total), desc=f"处理中 {i+1}/{total}...")

                    # 单张完成进度
                    progress((i + 1) / total, desc=f"完成 {i+1}/{total}")
                except Exception as e:
                    return f"处理第{i+1}张时出错: {e}", []

            return f"批量放大完成，成功处理 {total} 张图片", results
        except Exception as e:
            print(f"Error in batch upscale: {str(e)}")
            import traceback
            traceback.print_exc()
            return f"批量处理失败：{str(e)}", []

def Tab_image_upscale_batch():
    interface = BatchUpscaleInterface()
    
    # 放大模型选项
    upscale_model_choices = interface.upscale_models
    default_upscale_model = interface.default_upscale_model if interface.default_upscale_model and interface.default_upscale_model in upscale_model_choices else (upscale_model_choices[0] if upscale_model_choices else "RealESRGAN_x4plus.pth")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 批量图片放大设置")
            
            # 文件夹路径输入
            folder_path_input = gr.Textbox(
                label="图片文件夹路径", 
                placeholder="请输入包含图片的文件夹路径，例如：D:\\images",
                value=""
            )
            
            # 文件数量显示
            file_count_display = gr.Textbox(
                label="文件夹中的图片数量", 
                interactive=False,
                value="0"
            )
            
            # 放大模型选择
            upscale_model_dropdown = gr.Dropdown(
                label="放大模型选择", 
                choices=upscale_model_choices, 
                value=default_upscale_model, 
                interactive=bool(upscale_model_choices)
            )
            
            # 输出文件名前缀
            output_prefix_input = gr.Textbox(
                label="输出文件名前缀", 
                value="upscaled",
                placeholder="输出文件的前缀名称"
            )
            
            # 统计文件数量的函数
            def update_file_count(folder_path):
                count = interface.count_images_in_folder(folder_path)
                return str(count)
            
            # 当文件夹路径改变时，自动更新文件数量
            folder_path_input.change(
                fn=update_file_count,
                inputs=[folder_path_input],
                outputs=[file_count_display]
            )
        
        with gr.Column(scale=2):
            # 开始批量处理按钮
            batch_process_btn = gr.Button("开始批量放大", variant="primary", size="lg")
            
            # 状态信息显示
            status_text = gr.Textbox(label="处理状态", interactive=False, lines=3)
            
            # 结果图片展示
            result_gallery = gr.Gallery(
                label="处理结果预览", 
                columns=3, 
                show_label=True, 
                elem_id="batch_result_gallery"
            )
    
    # 绑定批量处理事件
    batch_process_btn.click(
        fn=interface.batch_upscale_images,
        inputs=[folder_path_input, upscale_model_dropdown, output_prefix_input],
        outputs=[status_text, result_gallery]
    )
    
    # 输出文件夹图片预览
    from config.config import get_output_path
    output_dir = get_output_path()
    
    def get_output_images():
        if not os.path.exists(output_dir):
            return []
        images = []
        for file in os.listdir(output_dir):
            if file.lower().endswith(('.png', '.jpg', '.jpeg')):
                images.append(os.path.join(output_dir, file))
        images.sort(key=lambda x: os.path.getmtime(x), reverse=True)
        return images
    
    with gr.Row():
        refresh_btn = gr.Button("刷新输出预览")
    with gr.Row():
        output_gallery = gr.Gallery(label="输出文件夹图片", columns=4, show_label=True, elem_id="output_gallery")
    
    refresh_btn.click(
        fn=get_output_images,
        outputs=[output_gallery]
    )